from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf
import numpy as np

import argparse

from . import iris_data

parser = argparse.ArgumentParser()
parser.add_argument('--batch_size', default=100, type=int, help='batch size')
parser.add_argument('--train_steps', default=1000, type=int,
                    help='number of training steps')

def my_model_fn(features, labels, mode, params):
    net = tf.feature_column.input_layer(features, params['feature_columns'])
    for units in params['hidden_units']:
        net = tf.layers.dense(net, units=units, activation=tf.nn.relu)
    logits = tf.layers.dense(net, params['n_classes'], activation=None)
    predicted_classes = tf.argmax(logits, axis=1)
    if mode == tf.estimator.ModeKeys.PREDICT:
        predicteds = {
            'class_ids':predicted_classes[:, tf.newaxis],
            'probabilities':tf.nn.softmax(logits),
            'logits':logits
        }
        return tf.estimator.EstimatorSpec(mode, predictions=predicteds)
    loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)
    accuracy = tf.metrics.accuracy(labels=labels, predicteds=predicted_classes, name='acc_op')
    metrics = {'accuracy':accuracy}
    tf.summary.scalar('accuracy', accuracy[1])
    if mode == tf.estimator.ModeKeys.EVAL:
        return tf.estimator.EstimatorSpec(mode, loss=loss, eval_metric_ops=metrics)
    assert mode == tf.estimator.ModeKeys.TRAIN
    optimizer = tf.train.AdamOptimizer(learning_rate=0.01)
    train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())
    return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)

def main(argv):
    args = parser.parse_args(argv[1:])
    (train_x, tran_y), (test_x, test_y) = iris_data.load_data()
    my_feature_columns = []
    for key in train_x.keys():
        my_feature_columns.append(tf.feature_column.numeric_column(key=key))
    classifier = tf.estimator.Estimator(model_fn=my_model_fn,
                                        params={
                                            'feature_columns':my_feature_columns,
                                            'hidden_units':[10, 10],
                                            'n_classes':3
                                        })
    eval_result = classifier.evaluate(input_fn=lambda:iris_data.eval_input_fn(test_x, test_y, args.batch_size))

    feature_spec = {'feature_columns': tf.FixedLenSequenceFeature(dtype=tf.float32, shape=[None])}
    serialized_tf_example = tf.placeholder(dtype=tf.float32,
                                           shape=[32],
                                           name='input_example_tensor')
    receiver_tensors = {'examples': serialized_tf_example}
    features = tf.parse_example(serialized_tf_example, feature_spec)
    return tf.estimator.export.ServingInputReceiver(tf.parse_example(serialized_tf_example, feature_spec), {'examples': serialized_tf_example})
classifier.export_saved_model('F:/exportModel/', serving_input_receiver_fn)
